#include <iostream>
#include <string>
#include <thread>
#include <vector>
#include <random>
#include <chrono>
#include "DatabaseFactory.h"

using Key = std::string;
using Value = int;

// 生成随机键
std::string generateRandomKey(std::mt19937& gen, std::uniform_int_distribution<>& keyDist) {
    return std::to_string(keyDist(gen));
}

// 压力测试函数
void stressTest(std::shared_ptr<KVDatabase<Key, Value>> db, int operations, double insertProbability) {
    std::random_device rd;
    std::mt19937 gen(rd());
    std::uniform_int_distribution<> keyDist(0, 9999); // 键范围 "0" 到 "999"
    std::uniform_int_distribution<> valueDist(0, 1000000); // 随机值范围
    std::uniform_real_distribution<> opDist(0.0, 1.0); // 操作概率分布

    for (int i = 0; i < operations; ++i) {
        std::string key = generateRandomKey(gen, keyDist);
        if (opDist(gen) < insertProbability) { // 插入操作
            Value value = valueDist(gen);
            db->put(key, value);
        } else { // 查询操作
            Value value;
            db->get(key, value);
        }
    }
}

int main() {
    using namespace std::chrono;

    const int operations = 1000000; // 每次测试的操作次数
    const double insertProbability = 0.1; // 插入操作的概率
    const int threadCounts[] = { 1, 2, 4 }; // 测试线程数


    for (int threadCount : threadCounts) {
        // 创建哈希表实现，并转换为 shared_ptr
        auto db = std::shared_ptr<KVDatabase<Key, Value>>(DatabaseFactory<Key, Value>::create(DatabaseType::HashMap));

        std::cout << "Testing with " << threadCount << " thread(s)..." << std::endl;

        auto start = high_resolution_clock::now();

        // 创建线程并执行压力测试
        std::vector<std::thread> threads;
        for (int i = 0; i < threadCount; ++i) {
            threads.emplace_back(stressTest, db, operations / threadCount, insertProbability);
        }

        // 等待所有线程完成
        for (auto& t : threads) {
            t.join();
        }

        auto end = high_resolution_clock::now();
        auto duration = duration_cast<milliseconds>(end - start).count();

        // 计算 QPS
        int qps = (operations * 1000.0) / duration;
        std::cout << "Completed " << operations << " operations in " << duration << " ms. QPS: " << qps << std::endl;
    }

    for (int threadCount : threadCounts) {
        // 创建哈希表实现，并转换为 shared_ptr
        auto db = std::shared_ptr<KVDatabase<Key, Value>>(DatabaseFactory<Key, Value>::create(DatabaseType::StripedLockingKVDB));

        std::cout << "Testing with " << threadCount << " thread(s)..." << std::endl;

        auto start = high_resolution_clock::now();

        // 创建线程并执行压力测试
        std::vector<std::thread> threads;
        for (int i = 0; i < threadCount; ++i) {
            threads.emplace_back(stressTest, db, operations / threadCount, insertProbability);
        }

        // 等待所有线程完成
        for (auto& t : threads) {
            t.join();
        }

        auto end = high_resolution_clock::now();
        auto duration = duration_cast<milliseconds>(end - start).count();

        // 计算 QPS
        int qps = (operations * 1000.0) / duration;
        std::cout << "Completed " << operations << " operations in " << duration << " ms. QPS: " << qps << std::endl;
    }


    return 0;
}